package cn.bugstack.chatglm.executor.aigc;

import cn.bugstack.chatglm.IOpenAiApi;
import cn.bugstack.chatglm.executor.Executor;
import cn.bugstack.chatglm.model.*;
import cn.bugstack.chatglm.session.Configuration;
import com.alibaba.fastjson.JSON;
import okhttp3.*;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;

import java.io.IOException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;

public class GLMExecutor implements Executor {
    private Configuration configuration;
    private EventSource.Factory factory;
    public GLMExecutor(Configuration configuration){
        this.configuration = configuration;
        factory = configuration.createRequestFactory();
    }
    @Override
    public ChatCompletionResponse chatCompletion(ChatCompletionRequest chatCompletionRequest) throws IOException {
        chatCompletionRequest.setStream(false);

        //新建request
        // 构建请求信息

        Request request = new Request.Builder()
                .url(configuration.getApiHost().concat(IOpenAiApi.v4))
                .post(RequestBody.create(MediaType.parse(Configuration.JSON_CONTENT_TYPE), JSON.toJSONString(chatCompletionRequest)))
                .addHeader("Content-Type", Configuration.JSON_CONTENT_TYPE) // 添加Content-Type头
                .build();
        OkHttpClient okHttpClient = configuration.getOkHttpClient();

        Response response = okHttpClient.newCall(request).execute();
        return JSON.parseObject(response.body().string(), ChatCompletionResponse.class);
    }

    @Override
    public EventSource completions(ChatCompletionRequest chatCompletionRequest, EventSourceListener eventSourceListener) throws IOException {
        if (!chatCompletionRequest.getStream()){
            throw new RuntimeException("非法输入字段stream是false");
        }

        //新建request
        // 构建请求信息
        Request request = new Request.Builder()
                .url(configuration.getApiHost().concat(IOpenAiApi.v4))
                .post(RequestBody.create(MediaType.parse(Configuration.JSON_CONTENT_TYPE), JSON.toJSONString(chatCompletionRequest)))
                .addHeader("Content-Type", Configuration.JSON_CONTENT_TYPE) // 添加Content-Type头
                .build();

        return factory.newEventSource(request, eventSourceListener);
    }

    @Override
    public ImageCompletionResponse genImages(ImageCompletionRequest imageCompletionRequest) throws IOException {
        Request request = new Request.Builder()
                .url(configuration.getApiHost().concat(IOpenAiApi.v4Image))
                .post(RequestBody.create(MediaType.parse(Configuration.JSON_CONTENT_TYPE), JSON.toJSONString(imageCompletionRequest)))
                .addHeader("Content-Type", Configuration.JSON_CONTENT_TYPE) // 添加Content-Type头
                .build();
        OkHttpClient okHttpClient = configuration.getOkHttpClient();

        Response response = okHttpClient.newCall(request).execute();
        return JSON.parseObject(response.body().string(), ImageCompletionResponse.class);
    }

    @Override
    public VideoCompletionResponse genVideos(VideoCompletionRequest videoCompletionRequest) throws Exception {
        // 1. 发起视频生成请求，生成视频ID
        String videoId = initiateVideoGeneration(videoCompletionRequest);

        // 2. 轮询视频生成进度
        ExecutorService executorService = Executors.newSingleThreadExecutor();
        Future<VideoCompletionResponse> resultFuture = executorService.submit(() -> pollForResult(videoId));

        try {
            // Wait for the result (can set timeout here if needed)
            VideoCompletionResponse video = resultFuture.get();
            return video;
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            executorService.shutdown();
        }
        return null;
    }

    private String initiateVideoGeneration(VideoCompletionRequest videoCompletionRequest) throws Exception {
        Request request = new Request.Builder()
                .url(configuration.getApiHost().concat(IOpenAiApi.VIDEO))
                .post(RequestBody.create(MediaType.parse(Configuration.JSON_CONTENT_TYPE), JSON.toJSONString(videoCompletionRequest)))
                .addHeader("Content-Type", Configuration.JSON_CONTENT_TYPE) // 添加Content-Type头
                .build();
        OkHttpClient okHttpClient = configuration.getOkHttpClient();

        Response response = okHttpClient.newCall(request).execute();
        VideoProcessResponse videoProcessResponse = JSON.parseObject(response.body().string(), VideoProcessResponse.class);
        return videoProcessResponse.getId();
    }

    private  VideoCompletionResponse pollForResult(String requestId) throws Exception {
        while (true) {
            Request request = new Request.Builder()
                    .url(configuration.getApiHost().concat(IOpenAiApi.GET_VIDEO_RES).concat(requestId))
                    .get()
                    .addHeader("Content-Type", Configuration.JSON_CONTENT_TYPE) // 添加Content-Type头
                    .build();
            OkHttpClient okHttpClient = configuration.getOkHttpClient();

            Response response = okHttpClient.newCall(request).execute();
            if (response.code() == 200) {
                // Assume response JSON has {"status":"success", "videoUrl":"<URL>"} when complete
                String responseBody = response.body().string();
                if (isSuccess(responseBody)) {
                    return extractVideo(responseBody);
                }
            }

            // Wait for 1 second before polling again
            Thread.sleep(1000);
        }
    }

    private  boolean isSuccess(String responseBody) {
        VideoCompletionResponse videoCompletionResponse = JSON.parseObject(responseBody, VideoCompletionResponse.class);
        if (videoCompletionResponse.getTask_status().equals("SUCCESS")){
            return true;
        }
        return false;
    }

    private VideoCompletionResponse extractVideo(String responseBody) {
        VideoCompletionResponse videoCompletionResponse = JSON.parseObject(responseBody, VideoCompletionResponse.class);
        return videoCompletionResponse;
    }
}
